package com.ai.util;

import com.fasterxml.jackson.databind.ObjectMapper;
import jakarta.annotation.PostConstruct;
import org.springframework.stereotype.Service;
import org.vosk.LibVosk;
import org.vosk.LogLevel;
import org.vosk.Model;
import org.vosk.Recognizer;
import javax.sound.sampled.AudioFormat;
import javax.sound.sampled.AudioInputStream;
import javax.sound.sampled.AudioSystem;
import java.io.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicReference;

/**
 * 语音转文字
 *
 * @author houwenpeng
 */
@Service
public class VoiceRecognition {

    // TODO 重量资源加载一次
    private Model model;
    private static final String MODEL_PATH = System.getProperty("user.dir") + "/src/main/resources/model/vosk-model-cn-0.22";
    @PostConstruct
    public void init() {
        LibVosk.setLogLevel(LogLevel.INFO);
        try {
            loadModel();  // 在应用启动时加载模型
        } catch (Exception e) {
            e.printStackTrace();
        }
    }
    private void loadModel() throws Exception {
        if (model == null) {
            model = new Model(MODEL_PATH);
        }
    }


    /**
     * 将音频转换为文本
     */
    public String convertSpeechToText(File file) {
        AtomicReference<String> resultText = new AtomicReference<>("");
        ExecutorService executorService = Executors.newSingleThreadExecutor(); // 使用单线程池，避免并发问题

        Future<String> future = executorService.submit(() -> {
            try (InputStream ais = AudioSystem.getAudioInputStream(new BufferedInputStream(new FileInputStream(file)));
                 Recognizer recognizer = new Recognizer(model, getSampleRate(file))) {

                recognizer.setWords(true); // 开启词信息的详细输出

                int bytes;
                byte[] b = new byte[8192]; // 增大缓存，提高读写效率
                while ((bytes = ais.read(b)) >= 0) {
                    recognizer.acceptWaveForm(b, bytes); // 分批次传递数据给 Vosk
                }

                String finalResult = recognizer.getFinalResult();

                // 解析 JSON 返回文本
                ObjectMapper objectMapper = new ObjectMapper();
                resultText.set(objectMapper.readTree(finalResult).get("text").asText());

            } catch (Exception e) {
                e.printStackTrace();
            }
            String s = resultText.get();
            System.out.println("音频结果:"+ s);
            return s;
        });

        try {
            return future.get(60, TimeUnit.SECONDS); // 等待结果，最多等待 60 秒
        } catch (TimeoutException e) {
            System.err.println("识别超时！");
            future.cancel(true);
        } catch (InterruptedException | ExecutionException e) {
            e.printStackTrace();
        } finally {
            executorService.shutdown();
        }
        String tests = resultText.get();
        System.out.println("音频结果:"+ tests);
        return tests;
    }

    /**
     * 获取音频的采样率，采样率为音频采样率 * 声道数
     */
    private float getSampleRate(File file) throws Exception {
        AudioInputStream audioInputStream = AudioSystem.getAudioInputStream(file);
        AudioFormat format = audioInputStream.getFormat();
        float sampleRate = format.getSampleRate();
        int channels = format.getChannels();
        return sampleRate * channels;
    }
}
